from solvers.solverbase import SolverBase
from solvers.docarmo.docarmobase import DoCarmoBase
from solvers.stabilizations import supg
from dolfin import *
import numpy

class GaussLegendre4(SolverBase,DoCarmoBase):
    '2-symplectic method rewritten without increments. Nonlinear-supg stabilized.'
    def __init__(self,p=1):
        SolverBase.__init__(self,p)
        DoCarmoBase.__init__(self)

    def solve(self,problem):
        mesh = problem.domain

        p = self.p
        V = FunctionSpace(mesh,'CG',p)
        VV = VectorFunctionSpace(mesh,'CG',p)
        DG = FunctionSpace(mesh,'DG',0)
        DGV = VectorFunctionSpace(mesh,'DG',0)

        #-----------------------------------------------------------------------------------------------------------------------------
        def projection(v,phi,phi_,p_,dt,cache):
            p = TrialFunction(DGV)
            phiT = TestFunction(DGV)

            if 'aVV' not in cache.keys():
                a = inner(p,phiT)*dx
                cache['aVV'] = assemble(a)

            w = conditional(
                        gt(sqrt(inner(nabla_grad(phi),nabla_grad(phi))),1E-16),\
                        (phi-phi_+dt*inner(v,nabla_grad(phi)))*nabla_grad(phi)/inner(nabla_grad(phi),nabla_grad(phi)),\
                        Constant((0.,0.))\
                           ) 

            l = inner(w,phiT)*dx
            b = assemble(l)
            solve(cache['aVV'],p_.vector(),b)
            
        def get_sigma(v,phi,phi_,p_,dt,sigma_,cache): # velocity,current sol,previous level solution,projection,coefs(output) 
            trial = TrialFunction(DG)
            psi = TestFunction(DG)

            ctrial = TrialFunction(V)
            cpsi = TestFunction(V)
            
            norm_v = Function(V)
            norm_p = Function(DG)
            resF = Function(DG)
            
            if 'aV' not in cache.keys():
                a =  trial*psi*dx
                cache['aV'] = assemble(a)
            
            if 'caV' not in cache.keys():
                ca = ctrial*cpsi*dx
                cache['caV'] = assemble(ca)

            #-----
            p_norm_l = sqrt(inner(p_,p_))*psi*dx
            b1 = assemble(p_norm_l)
            solve(cache['aV'],norm_p.vector(),b1)               # get the norm of p_(discon)
            #-----
            v_norm_l = sqrt(inner(v,v))*cpsi*dx
            b2 = assemble(v_norm_l)
            solve(cache['caV'],norm_v.vector(),b2)              # get the norm of v (continuous) ! 
            #-----
            resl = (phi-phi_+dt*inner(v,nabla_grad(phi)))*psi*dx
            b3 = assemble(resl)
            solve(cache['aV'],resF.vector(),b3)                 # get the residuum
            
            for i in xrange(len(sigma_.vector().array())):
                c = Cell(mesh,i)
                x = c.midpoint()

                v_mag = norm_v(x.x(),x.y())
                p_mag = norm_p(x.x(),x.y())
                residuum = resF(x.x(),x.y())
                
                if abs(residuum) < 1E-16:
                    sigma_.vector()[i] = 0
                else:
                    alfa = 1.
                    if p_mag >= v_mag:
                        sigma_.vector()[i] = 0
                    else:
                        sigma_.vector()[i] = v_mag*p_mag - alfa*p_mag**2 
        #----------------------------------------------------------------------------------------------------------------------------------
 
        phi_ = interpolate(problem.phi_,V)
        phiP = Function(V)                  # solution to the auxiliary equation
        phiC = Function(V)                  # solution to the true equation
        phi = TrialFunction(V)          
        varphi = TestFunction(V)

        t = problem.t          
        dt = Constant(SolverBase._get_time_step(self,problem,VV,t))
        T = problem.T

        v = SolverBase._get_velocity(self,problem,VV,t)
        
        tau = Function(DG)
        tau.vector()[:] = supg.get_tau(v,supg.diffusivity)  # stab. parameters

        p_ = Function(DGV) 
        sigma_p = Function(DG) 
        sigma_m = Function(DG)
        pi_p = Function(DG)
        pi_m = Function(DG) 
        
        pert = conditional(
                           gt(sqrt(inner(v,v)),1E-16),\
                           Constant(1.0),\
                           Constant(0.) \
                           ) # a way to realize swith on |v| == 0.

        aP = phi*varphi*dx\
            + tau*pert*phi*inner(v,nabla_grad(varphi))*dx\
            + 0.5*dt*sqrt(3.)/6.*sigma_m*tau*inner(nabla_grad(phi),nabla_grad(varphi))*dx\
            - 0.5*dt*sqrt(3.)/6.*sigma_p*tau*inner(nabla_grad(phi),nabla_grad(varphi))*dx

        lP = phi_*varphi*dx\
             + dt*sqrt(3.)/6.*inner(v,nabla_grad(phiC))*varphi*dx \
             - dt*sqrt(3.)/6.*inner(v,nabla_grad(phi_))*varphi*dx\
             + tau*pert*phi_*inner(v,nabla_grad(varphi))*dx \
             + dt*sqrt(3.)/6.*tau*pert*inner(dot(nabla_grad(phiC),outer(v,v)),nabla_grad(varphi))*dx \
             - dt*sqrt(3.)/6.*tau*pert*inner(dot(nabla_grad(phi_),outer(v,v)),nabla_grad(varphi))*dx\
             -0.5*dt*1./2.*sigma_m*tau*inner(nabla_grad(phi_),nabla_grad(varphi))*dx\
             -0.5*dt*(1./2.-sqrt(3.)/6.)*sigma_m*tau*inner(nabla_grad(phiC),nabla_grad(varphi))*dx\
             +0.5*dt*(1./2.-2.*sqrt(3.)/6.)*sigma_p*tau*inner(nabla_grad(phi_),nabla_grad(varphi))*dx\
             +0.5*dt*(1./2.+sqrt(3.)/6.)*sigma_p*tau*inner(nabla_grad(phiC),nabla_grad(varphi))*dx
#-----
        aC = phi*varphi*dx + 0.5*dt*inner(v,nabla_grad(phi))*varphi*dx\
             +tau*pert*phi*inner(v,nabla_grad(varphi))*dx\
             +0.5*dt*tau*pert*inner(dot(nabla_grad(phi),outer(v,v)),nabla_grad(varphi))*dx\
             +0.5*dt*(1./2.-sqrt(3.)/6.)*pi_m*tau*inner(nabla_grad(phi),nabla_grad(varphi))*dx\
             +0.5*dt*(1./2.+sqrt(3.)/6.)*pi_p*tau*inner(nabla_grad(phi),nabla_grad(varphi))*dx
    
        lC = phi_*varphi*dx + tau*pert*phi_*inner(v,nabla_grad(varphi))*dx\
             -dt*(0.5-sqrt(3.)/6.)*inner(v,nabla_grad(phi_))*varphi*dx\
             -dt*(0.5-sqrt(3.)/6.)*tau*pert*inner(dot(nabla_grad(phi_),outer(v,v)),nabla_grad(varphi))*dx\
             -dt*sqrt(3.)/6.*inner(v,nabla_grad(phiP))*varphi*dx\
             -dt*sqrt(3.)/6.*tau*pert*inner(dot(nabla_grad(phiP),outer(v,v)),nabla_grad(varphi))*dx\
             -0.5*dt*1./2.*pi_m*tau*inner(nabla_grad(phi_),nabla_grad(varphi))*dx\
             -0.5*dt*sqrt(3.)/6.*pi_m*tau*inner(nabla_grad(phiP),nabla_grad(varphi))*dx\
             -0.5*dt*(1./2.-2.*sqrt(3.)/6.)*pi_p*tau*inner(nabla_grad(phi_),nabla_grad(varphi))*dx\
             -0.5*dt*sqrt(3.)/6.*pi_p*tau*inner(nabla_grad(phiP),nabla_grad(varphi))*dx

        self.update(problem,phi_,t,float(dt),0)  # store the solution at time 0

        phiC.assign(phi_)   # initial guess
        
        phiC_ = Function(V) # so that error could be calculated
        phiC_.assign(phi_)

        phi = Function(V)
        phi_k = Function(V) # iterations for 2nd eq
        phi_m = Function(V) # iterations for 2nd eq
   
        t = 0

        count = 0
        cache = {}
        maxIter = 10
        tolG = 1E-8            # tolerance for gauss loop
        tolP = 1E-3            # tolerance for picard loop
        
        while t <= T:  # time loop
            t += float(dt)
            count += 1
            
            errorG = 1.
            iterationG = 0
           
            while errorG > tolG and iterationG < maxIter: # gauss4 loop
                iterationG += 1

                errorP1 = 1.
                iterationP1 = 0
                # guess for sigma_m 
                projection(v,0.5*phi_+(0.5-sqrt(3.)/6.)*phiC+sqrt(3.)/6.*phi_,phi_,p_,dt,cache)
                get_sigma(v,0.5*phi_+(0.5-sqrt(3.)/6.)*phiC+sqrt(3.)/6.*phi_,phi_,p_,float(dt),sigma_m,cache) 
                # guess for sigma_p
                projection(v,(0.5-2*sqrt(3.)/6.)*phi_+(0.5+sqrt(3.)/6.)*phiC+sqrt(3.)/6.*phi_,phi_,p_,dt,cache)
                get_sigma(v,(0.5-2*sqrt(3.)/6.)*phi_+(0.5+sqrt(3.)/6.)*phiC+sqrt(3.)/6.*phi_,phi_,p_,float(dt),sigma_p,cache)  
                while errorP1 > tolP and iterationP1 < maxIter: # picard loop for second eq.
                    iterationP1 += 1
                    AP = assemble(aP)
                    LP = assemble(lP)
                    solve(AP,phi.vector(),LP)  # solve true system, stabilized -> nonlinear
                    diffP1 = phi.vector().array() - phi_k.vector().array()
                    errorP1 = numpy.linalg.norm(diffP1,ord=numpy.Inf)
                    phi_k.assign(phi)
                    # guess for sigma_m 
                    projection(v,0.5*phi_+(0.5-sqrt(3.)/6.)*phiC+sqrt(3.)/6.*phi_k,phi_,p_,dt,cache)
                    get_sigma(v,0.5*phi_+(0.5-sqrt(3.)/6.)*phiC+sqrt(3.)/6.*phi_k,phi_,p_,float(dt),sigma_m,cache) 
                    # guess for sigma_p
                    projection(v,(0.5-2*sqrt(3.)/6.)*phi_+(0.5+sqrt(3.)/6.)*phiC+sqrt(3.)/6.*phi_k,phi_,p_,dt,cache)
                    get_sigma(v,(0.5-2*sqrt(3.)/6.)*phi_+(0.5+sqrt(3.)/6.)*phiC+sqrt(3.)/6.*phi_k,phi_,p_,float(dt),sigma_p,cache)  
                    #print '\tpicard: ',iterationP1,errorP1
                phiP.assign(phi)

                errorP2 = 1.
                iterationP2 = 0
                # guess for pi_m 
                projection(v,0.5*phi_+(0.5-sqrt(3.)/6.)*phiC+sqrt(3.)/6.*phiP,phi_,p_,dt,cache)
                get_sigma(v,0.5*phi_+(0.5-sqrt(3.)/6.)*phiC+sqrt(3.)/6.*phiP,phi_,p_,float(dt),pi_m,cache) 
                # guess for pi_p
                projection(v,(0.5-2*sqrt(3.)/6.)*phi_+(0.5+sqrt(3.)/6.)*phiC+sqrt(3.)/6.*phiP,phi_,p_,dt,cache)
                get_sigma(v,(0.5-2*sqrt(3.)/6.)*phi_+(0.5+sqrt(3.)/6.)*phiC+sqrt(3.)/6.*phiP,phi_,p_,float(dt),pi_p,cache)  
                while errorP2 > tolP and iterationP2 < maxIter: # picard loop for second eq.
                    iterationP2 += 1
                    AC = assemble(aC)
                    LC = assemble(lC)
                    solve(AC,phi.vector(),LC)  # solve true system, stabilized -> nonlinear
                    diffP2 = phi.vector().array() - phi_m.vector().array()
                    errorP2 = numpy.linalg.norm(diffP2,ord=numpy.Inf)
                    phi_m.assign(phi)
                    # guess for pi_m 
                    projection(v,0.5*phi_+(0.5-sqrt(3.)/6.)*phi_m+sqrt(3.)/6.*phiP,phi_,p_,dt,cache)
                    get_sigma(v,0.5*phi_+(0.5-sqrt(3.)/6.)*phi_m+sqrt(3.)/6.*phiP,phi_,p_,float(dt),pi_m,cache) 
                    # guess for pi_p
                    projection(v,(0.5-2*sqrt(3.)/6.)*phi_+(0.5+sqrt(3.)/6.)*phi_m+sqrt(3.)/6.*phiP,phi_,p_,dt,cache)
                    get_sigma(v,(0.5-2*sqrt(3.)/6.)*phi_+(0.5+sqrt(3.)/6.)*phi_m+sqrt(3.)/6.*phiP,phi_,p_,float(dt),pi_p,cache)  
                    #print '\tpicard: ',iterationP2,errorP2
                phiC.assign(phi)

                diffG = phiC.vector().array() - phiC_.vector().array()
                errorG = numpy.linalg.norm(diffG,ord=2)        # compute the error
                phiC_.assign(phiC)
                #print 'gauss: ',iterationG,errorG

            phi_.assign(phiC)
            self.update(problem,phi_,t,float(dt),count)

        problem.save_data(self.saveDir,str(self.CFL))
     
